#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
import os
import datetime
import time
from bs4 import BeautifulSoup
import sys

HOST = 'localhost'
PORT = 27017
DB_NAME = 'pmc'
COLLECTION_NAME = 'metalist'


def GetConfig():
    if os.path.exists("../config.ini"):
        import ConfigParser
        global HOST
        global PORT
        global DB_NAME
        global COLLECTION_NAME
        config = ConfigParser.ConfigParser()
        config.read("../config.ini")
        HOST = config.get("mongodb", "host")
        PORT = int(config.get("mongodb", "port"))
        DB_NAME = config.get("mongodb", "db")
        COLLECTION_NAME = config.get("mongodb", "collection")


def Get_PMCID_Manuscript(article_string):
    soup = BeautifulSoup(article_string, "xml")
    pmcid = None
    manuscriptid = None
    pmcid_node = soup.find("article-id", attrs={"pub-id-type": "pmc"})
    if pmcid_node != None:
        pmcid = pmcid_node.string
    manuscript_node = soup.find("article-id", attrs={"pub-id-type": "manuscript"})
    if manuscript_node != None:
        manuscriptid = manuscript_node.string
    if pmcid != None and pmcid.isdigit():
        return (pmcid + "\t" + str(manuscriptid))
    else:
        return None


# Use beautifulsoup to analyze PMCID lines
def Extract_PMC_IDs_From_XML(filename):
    infile = open(filename, "r")
    result = []
    article_string = ""
    is_string_part = False
    for line in infile:
        line = line.rstrip()
        if line.find("<article-meta>") != -1:
            is_string_part = True
        if is_string_part:
            article_string += line
        if line.find("</article-meta>") != -1:
            is_string_part = False
            pmcid_manuscriptid = Get_PMCID_Manuscript(article_string)
            if pmcid_manuscriptid != None:
                result.append(pmcid_manuscriptid)
            article_string = ""
    infile.close()
    return result


if __name__ == '__main__':
    os.chdir(sys.argv[1])

    print("Loading date, given date is ... "),
    time.sleep(1)
    dateString = sys.argv[1].split("_")[-1]
    dateObject = datetime.datetime.strptime(dateString, "%Y%m%d")
    print(dateString + ".")
    time.sleep(1)

    print("Loading config file ... "),
    time.sleep(1)
    GetConfig()
    print("complete.")
    time.sleep(1)

    print("Initiating mongodb connection ... "),
    time.sleep(1)
    connection = MongoClient(HOST, PORT)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]
    print("complete.")
    time.sleep(1)

    print("Loading PMC xml files ... "),
    time.sleep(1)
    filelist = os.listdir(".")
    print("complete.")
    time.sleep(1)

    print("Start anlyzing xml files ... ")
    time.sleep(1)
    for filename in filelist:
        if filename.endswith(".xml"):
            manuscript_count = 0
            print("Start analyzing " + filename + " ... "),
            filepath = os.path.dirname(os.path.abspath(filename)).split("/")[-1] + "/" + filename
            ids = Extract_PMC_IDs_From_XML(filename)
            for id in ids:
                split_id = id.split("\t")
                pmcid = int(split_id[0])
                manuscriptid = split_id[1]
                if manuscriptid == "None":
                    collection.update_one({"pmcid": pmcid},
                                          {"$set": {"xml": filepath, "time_in": dateObject, "time_last": dateObject}},
                                          upsert=True)
                else:
                    manuscript_count += 1
                    collection.update_one({"pmcid": pmcid}, {
                        "$set": {"xml": filepath, "time_in": dateObject, "time_last": dateObject,
                                 "manuscript_id": manuscriptid}}, upsert=True)
            print("complete. Find " + str(len(ids)) + " PMCIDs, " + str(manuscript_count) + " manuscripts.")
    connection.close()
    print("All xml files analyzed.")
